Skip to content

fix(Tokenizers): avoid quadratic regex matching for added tokens without whitespace padding - #386

Open
shubhransh-gupta wants to merge 1 commit into
huggingface:mainfrom
shubhransh-gupta:fix/added-tokens-regex-perf
Open

fix(Tokenizers): avoid quadratic regex matching for added tokens without whitespace padding#386
shubhransh-gupta wants to merge 1 commit into
huggingface:mainfrom
shubhransh-gupta:fix/added-tokens-regex-perf

Conversation

@shubhransh-gupta

Copy link
Copy Markdown

Fixes #383.

Summary

In PreTrainedTokenizer, addedTokensRegex was constructed by wrapping every added token in a capturing group (\(token)) even when neither lstrip nor rstrip was requested.

For tokenizers with large added token vocabularies (e.g. Google Gemma-3 with 6,415 added tokens), this caused an $O(N^2)$ quadratic slowdown in Foundation's ICU regular expression engine due to capturing group table allocations, slowing down single newline \n encoding from ~0.02 ms to ~29 ms (~1.1s for a 40-line prompt on device).

Changes

  1. Non-capturing groups by default: When $0.prefix and $0.suffix are both false, compile the token using a non-capturing group (?:\(token)) instead of (\(token)).
  2. Preserve whitespace stripping: Added tokens with lstrip or rstrip still retain their capturing group so split(by: NSRegularExpression) continues to isolate and strip surrounding whitespace as expected.
  3. Empty check guard: Avoid compiling an empty regex pattern when unwrappedAddedTokens is empty.
  4. Regression test: Added addedTokensRegexPerformanceAndCorrectness in Tests/TokenizersTests/TokenizerTests.swift.

Verification

  • Executed swift test across all targets (50/50 tests passing).
  • Validated code style with swift-format lint --strict --recursive Sources Tests (clean, 0 warnings).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Tokenizer.encode is quadratic in the number of added tokens (~29 ms per newline with Gemma-3)

1 participant